import { api } from "../../scripts/api.js";
import { app } from "../../scripts/app.js";

const dir = "workflows";
let allFiles = []; // [{  workflow_id, path }]

const isApiJson = (data) => {
  return (
    typeof data === "object" && Object.values(data).every((v) => v.class_type)
  );
};

const matchWorkflow = (file) => {
  return allFiles.find((item) => {
    if (file.startsWith(dir)) {
      return `${dir}/${item.path}` === file;
    } else {
      return item.path === file;
    }
  });
};

// 重写 loadApiJson函数
const rewriteLoadApiJson = () => {
  const originLoadGraphData = app.loadGraphData;

  app.loadGraphData = async (
    graphData,
    clean,
    restore_view,
    workflow,
    options
  ) => {
    if (
      workflow &&
      workflow.content &&
      isApiJson(JSON.parse(workflow.content))
    ) {
      app.graph.clear();
      app.loadApiJson(JSON.parse(workflow.content), workflow.filename);
      return;
    } else if (isApiJson(graphData)) {
      return;
    }
    return originLoadGraphData.apply(app, [
      graphData,
      clean,
      restore_view,
      workflow,
      options,
    ]);
  };
};

// 重写保存API
const rewriteStoreUserData = () => {
  const originStoreUserData = api.storeUserData;

  api.storeUserData = async (file, data, options) => {
    console.log("file", file, data);
    try {
      if (options?.overwrite) {
        // 保存
        const info = matchWorkflow(file);
        const newData = await app.graphToPrompt(app.graph);
        const res = await originStoreUserData.apply(api, [
          info && info.workflow_id ? `${dir}/${info.workflow_id}` : file,
          JSON.stringify(newData),
          options,
        ]);
        await loadWorkflows();
        return res;
      } else {
        // 另存
        const newData = await app.graphToPrompt(app.graph);
        const res = await originStoreUserData.apply(api, [
          file,
          JSON.stringify(newData),
          {
            ...options,
            stringify: false,
          },
        ]);
        await loadWorkflows();
        return res;
      }
    } catch {}
  };
};

// 重写删除API
const rewriteDeleteUserData = () => {
  const originDeleteUserData = api.deleteUserData;
  api.deleteUserData = async (file) => {
    const info = matchWorkflow(file);

    const res = await originDeleteUserData.apply(api, [
      info && info.workflow_id ? `${dir}/${info.workflow_id}` : file,
    ]);
    await loadWorkflows();
    return res;
  };
};

// 重写moveUserData
const rewriteMoveUserData = () => {
  const originMoveUserData = api.moveUserData;

  api.moveUserData = async (source, dest, options) => {
    const info = allFiles.find((item) => `${dir}/${item.path}` === source);

    console.log("source, dest, options", source, dest, options);
    const res = await originMoveUserData.apply(api, [
      info && info.workflow_id ? info.workflow_id : source,
      dest,
      options,
    ]);
    await loadWorkflows();
    return res;
  };
};

// 重写获取用户数据
const rewriteGetUserData = () => {
  const originGetUserData = api.getUserData;

  api.getUserData = async (file, options) => {
    const info = matchWorkflow(file);
    return originGetUserData.apply(api, [
      info && info.workflow_id ? `${dir}/${info.workflow_id}` : file,
      options,
    ]);
  };
};

// 加载云服务上全部workflow
const loadWorkflows = async () => {
  const res = await api.listUserDataFullInfo(dir);
  allFiles = Array.isArray(res) ? res : [];
};

// 重写API
const rewriteApis = () => {
  rewriteStoreUserData();
  rewriteDeleteUserData();
  rewriteGetUserData();
  rewriteMoveUserData();
};

const getHashParams = (key) => {
  const hash = window.location.hash.substring(1); // 去掉开头的 '#'
  const params = {};

  if (hash) {
    const pairs = hash.split("&");
    pairs.forEach((pair) => {
      const [key, value] = pair.split("=");
      if (key) {
        params[decodeURIComponent(key)] = value
          ? decodeURIComponent(value)
          : "";
      }
    });
  }

  return params[key];
};

const ext = {
  name: "Comfy.TransferWorkflowData",
  async setup() {
    rewriteLoadApiJson();
    rewriteApis();
    await loadWorkflows();
    const workflow_id = parseInt(getHashParams("workflow_id"));
    if (workflow_id) {
      const file_name =
        allFiles.find((item) => item.workflow_id === workflow_id)?.path || "";
      if (file_name) {
        try {
          const sid = sessionStorage.getItem("clientId");
          const res = await api.getUserData(file_name);
          let data = await res.json();

          if (isApiJson(data)) {
            await app.loadApiJson(data, file_name);
            const newData = await app.graphToPrompt(app.graph);
            data = newData.workflow;
          }

          if (sid) {
            const fileInfo = allFiles.find(
              (item) => item.workflow_id === workflow_id
            );
            const new_file_name = fileInfo?.status === 13 ? "" : file_name;
            sessionStorage.setItem(
              "Comfy.OpenWorkflowsPaths:" + sid,
              JSON.stringify(["workflows/" + new_file_name])
            );
            sessionStorage.setItem(
              "Comfy.PreviousWorkflow:" + sid,
              new_file_name
            );
            sessionStorage.setItem("workflow:" + sid, JSON.stringify(data));

            localStorage.setItem(
              "Comfy.OpenWorkflowsPaths",
              JSON.stringify(["workflows/" + new_file_name])
            );
            localStorage.setItem("Comfy.PreviousWorkflow", new_file_name);
            localStorage.setItem("workflow", JSON.stringify(data));

            history.replaceState(
              null,
              "",
              window.location.pathname + window.location.search
            );
          }
        } catch (error) {}
      }
    }
  },
};

app.registerExtension(ext);
